{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最小单元对齐模块 (MUAM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class MultiHeadCrossAttention(nn.Module):\n",
    "    def __init__(self, embed_dim, num_heads):\n",
    "        super(MultiHeadCrossAttention, self).__init__()\n",
    "        self.multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)\n",
    "        self.norm = nn.LayerNorm(embed_dim)\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(embed_dim, embed_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(embed_dim, embed_dim)\n",
    "        )\n",
    "\n",
    "    def forward(self, text_embeds, image_embeds):\n",
    "        # 交叉注意力机制, 将图像和文本特征对齐\n",
    "        attn_output, _ = self.multihead_attn(text_embeds, image_embeds, image_embeds)\n",
    "        text_embeds = text_embeds + attn_output  # 残差连接\n",
    "        text_embeds = self.norm(text_embeds)  # LayerNorm\n",
    "        text_embeds = self.mlp(text_embeds)  # MLP层\n",
    "        return text_embeds\n",
    "\n",
    "class MUAM(nn.Module):\n",
    "    def __init__(self, embed_dim, num_heads):\n",
    "        super(MUAM, self).__init__()\n",
    "        self.cross_attention = MultiHeadCrossAttention(embed_dim, num_heads)\n",
    "        self.fc = nn.Linear(embed_dim, 1)\n",
    "\n",
    "    def forward(self, text_embeds, image_embeds):\n",
    "        # 文本和图像特征对齐\n",
    "        aligned_text_embeds = self.cross_attention(text_embeds, image_embeds)\n",
    "\n",
    "        # 计算最小单元对齐度分数\n",
    "        similarity_matrix = torch.matmul(aligned_text_embeds, image_embeds.transpose(-1, -2))\n",
    "        alignment_scores = F.softmax(similarity_matrix, dim=-1)\n",
    "        alignment_scores = self.fc(alignment_scores).squeeze(-1)  # 压缩维度为[batch_size, num_image_patches]\n",
    "        return alignment_scores"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
